import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt  # 添加matplotlib用于绘图

# 检查CUDA是否可用，如果可用则使用GPU，否则使用CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 定义数据预处理方式
# transforms.Compose用于组合多个变换操作
transform = transforms.Compose([
    # 将PIL Image或numpy.ndarray转换为tensor格式
    transforms.ToTensor(),
    # 对数据进行标准化处理
    # 第一个参数是均值，第二个参数是标准差
    # 这些值是MNIST数据集的统计值
    transforms.Normalize((0.1307,), (0.3081,))
])

# 加载训练数据集
# root: 数据存储路径
# train=True: 表示加载训练集
# download=True: 如果数据不存在则自动下载
# transform: 应用之前定义的数据预处理
train_dataset = torchvision.datasets.MNIST(
    root='./data',      # 数据存储路径
    train=True,         # 加载训练集
    download=True,      # 自动下载数据
    transform=transform # 数据预处理
)

# 创建训练数据加载器
# batch_size: 每个批次包含64个样本
# shuffle=True: 每个epoch都打乱数据顺序
train_loader = torch.utils.data.DataLoader(
    train_dataset,      # 数据集
    batch_size=64,      # 批次大小
    shuffle=True        # 是否打乱数据
)

# 加载测试数据集
test_dataset = torchvision.datasets.MNIST(
    root='./data',       # 数据存储路径
    train=False,         # 加载测试集
    download=True,       # 自动下载数据
    transform=transform  # 数据预处理
)

# 创建测试数据加载器
test_loader = torch.utils.data.DataLoader(
    test_dataset,        # 数据集
    batch_size=1000,     # 批次大小
    shuffle=False        # 测试集不需要打乱
)

# 定义2层CNN网络
class Net(nn.Module):
    """
    2层CNN网络用于MNIST手写数字识别
    网络结构：
    1. 2个卷积层 + 池化层
    2. 1个全连接层 + Dropout
    3. 1个输出层
    """
    def __init__(self):
        """
        初始化网络结构
        """
        # 调用父类的初始化方法
        super(Net, self).__init__()
        
        # 第一个卷积层
        # in_channels: 输入通道数，MNIST是灰度图，所以是1
        # out_channels: 输出通道数，即卷积核的个数，这里是32个
        # kernel_size: 卷积核大小，这里是5x5
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=5)
        
        # 第二个卷积层
        # in_channels: 输入通道数，等于上一层的输出通道数，这里是32
        # out_channels: 输出通道数，这里是64个卷积核
        # kernel_size: 卷积核大小，仍然是5x5
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5)
        
        # Dropout层，用于防止过拟合
        # p=0.5: 每个神经元有50%的概率被暂时丢弃
        #大约一半（32个）的通道会被完全置零（整个4x4特征图变为全0）
        #剩余的通道会被放大2倍以保持期望值
        self.dropout = nn.Dropout2d(p=0.5)
        
        # 第一个全连接层
        # in_features: 输入特征数，这个值需要根据输入图像大小和网络结构计算得出
        # out_features: 输出特征数，这里设置为128
        self.fc1 = nn.Linear(in_features=1024, out_features=128)
        
        # 第二个全连接层（输出层）
        # in_features: 输入特征数，等于上一层的输出特征数，这里是128
        # out_features: 输出特征数，MNIST有10个类别（0-9），所以是10
        self.fc2 = nn.Linear(in_features=128, out_features=10)

    def forward(self, x):
        """
        前向传播函数
        :param x: 输入数据
        :return: 网络输出
        """
        # 第一个卷积层 + ReLU激活函数
        # F.relu: ReLU激活函数，将小于0的值置为0
        x = F.relu(self.conv1(x))
        
        # 第一个池化层
        # F.max_pool2d: 2维最大池化操作
        # kernel_size=2: 池化窗口大小为2x2
        x = F.max_pool2d(x, kernel_size=2)
        
        # 第二个卷积层 + ReLU激活函数
        x = F.relu(self.conv2(x))
        
        # 第二个池化层
        x = F.max_pool2d(x, kernel_size=2)
        
        # Dropout层
        x = self.dropout(x)
        
        # 展平操作，将多维特征图转换为一维向量
        # -1: 自动计算该维度的大小，保持总元素数量不变
        #1024：将后面的所有维度展平为1024个元素
        #x = x.view(-1, 1024)
        #更好的方案，x = x.view(x.shape[0], -1)，动态计算第二维度，增强代码的适应性
        x = x.view(x.shape[0], -1)

        # 第一个全连接层 + ReLU激活函数
        x = F.relu(self.fc1(x))
        
        # 第二个全连接层（输出层）
        x = self.fc2(x)
        
        # LogSoftmax激活函数，用于多分类任务的输出
        # dim=1: 在类别维度上进行softmax操作
        return F.log_softmax(x, dim=1)

# 创建网络实例
# Net(): 创建网络对象
# to(device): 将网络移动到指定设备（CPU或GPU）
model = Net().to(device)

# 定义损失函数
# nn.NLLLoss(): 负对数似然损失函数，适用于多分类问题
# 通常与log_softmax输出配合使用
criterion = nn.NLLLoss()

# 定义优化器
# optim.SGD: 随机梯度下降优化算法
# model.parameters(): 网络中所有可学习的参数
# lr=0.01: 学习率，控制参数更新的步长
# momentum=0.5: 动量参数，用于加速收敛
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)

# 添加列表用于存储训练和测试损失
train_losses = []
test_losses = []
train_counter = []
test_counter = []
def train(epoch):
    """
    训练函数
    :param epoch: 当前epoch数
    """
    # 将模型设置为训练模式
    # 启用dropout等训练时特有的操作
    model.train()
    
    # 遍历训练数据加载器中的每个批次
    for batch_idx, (data, target) in enumerate(train_loader):
        # 将数据和标签移动到指定设备上
        data, target = data.to(device), target.to(device)
        
        # 清零梯度缓存，防止梯度累积
        optimizer.zero_grad()
        
        # 前向传播，得到模型输出
        output = model(data)
        
        # 计算损失值
        loss = criterion(output, target)
        
        # 反向传播，计算梯度
        loss.backward()
        
        # 更新网络参数
        optimizer.step()
        
        # 每100个批次打印一次训练信息
        if batch_idx % 100 == 0:
            print(f'训练Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '
                  f'({100. * batch_idx / len(train_loader):.0f}%)]\t损失: {loss.item():.6f}')
            # 记录训练损失
            train_losses.append(loss.item())
            train_counter.append(batch_idx * len(data) + (epoch - 1) * len(train_loader.dataset))



def test():
    """
    测试函数
    """
    # 将模型设置为评估模式
    # 禁用dropout等训练时特有的操作
    model.eval()
    
    # 初始化测试损失和正确率
    test_loss = 0
    correct = 0
    
    # 不计算梯度，提高推理速度并减少内存消耗
    with torch.no_grad():
        # 遍历测试数据加载器中的每个批次
        for data, target in test_loader:
            # 将数据和标签移动到指定设备上
            data, target = data.to(device), target.to(device)
            
            # 前向传播，得到模型输出
            output = model(data)
            
            # 累加测试损失（取每个批次的平均损失）
            test_loss += criterion(output, target).item()
            
            # 获取预测结果（概率最大的类别）
            # output.max(1)返回(值, 索引)的元组，我们只需要索引
            pred = output.max(1, keepdim=True)[1]
            
            # 统计正确预测的样本数
            # pred.eq(target.view_as(pred))比较预测值和真实值是否相等
            # sum()统计正确预测的数量
            correct += pred.eq(target.view_as(pred)).sum().item()
    
    # 计算平均损失
    test_loss /= len(test_loader)

    # 记录测试损失
    test_losses.append(test_loss)
    
    # 打印测试结果
    print(f'\n测试集: 平均损失: {test_loss:.4f}, '
          f'准确率: {correct}/{len(test_loader.dataset)} '
          f'({100. * correct / len(test_loader.dataset):.2f}%)\n')

# 主训练循环
print("开始训练...")
# 在训练开始前进行一次测试
test()
test_counter = [i * len(train_loader.dataset) for i in range(11)]  # 11是因为我们有10个epoch，加上初始测试


# 训练10个epoch
for epoch in range(1, 11):
    # 调用训练函数
    train(epoch)
    # 每个epoch结束后进行测试
    test()

print("训练完成！")

# 绘制损失曲线
plt.figure(figsize=(10, 5))
plt.plot(train_counter, train_losses, color='blue', label='训练损失')
plt.scatter(test_counter, test_losses, color='red', label='测试损失')
plt.legend()
plt.xlabel('训练样本数')
plt.ylabel('损失')
plt.title('训练和测试损失曲线')
plt.show()